from __future__ import annotations

from metasim.constants import PhysicStateType
from metasim.scenario.objects import RigidObjCfg
from metasim.scenario.scenario import ScenarioCfg
from metasim.task.registry import register_task

from .rl_bench import RLBenchTask


@register_task(
    "rlbench.put_umbrella_in_umbrella_stand", "put_umbrella_in_umbrella_stand", "franka.put_umbrella_in_umbrella_stand"
)
class PutUmbrellaInUmbrellaStandTask(RLBenchTask):
    max_episode_steps = 200
    scenario = ScenarioCfg(
        objects=[
            RigidObjCfg(
                name="umbrella_visual",
                usd_path="roboverse_data/assets/rlbench/put_umbrella_in_umbrella_stand/umbrella_visual/usd/umbrella_visual.usd",
                physics=PhysicStateType.RIGIDBODY,
            ),
            RigidObjCfg(
                name="stand_visual",
                usd_path="roboverse_data/assets/rlbench/put_umbrella_in_umbrella_stand/stand_visual/usd/stand_visual.usd",
                physics=PhysicStateType.GEOM,
            ),
        ],
        robots=["franka"],
    )
    traj_filepath = "roboverse_data/trajs/rlbench/put_umbrella_in_umbrella_stand/v2/franka_v2.pkl.gz"
    # TODO: add checker
